from langchain_community.document_loaders import PyPDFLoader
import asyncio
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_ollama.embeddings import OllamaEmbeddings

# 替换 OpenAIEmbeddings 配置为 Ollama 本地模型
embeddings = OllamaEmbeddings(
    model="nomic-embed-text",  # 或其他支持的本地模型如 "mistral", "nomic-embed-text"
    base_url="http://192.168.43.158:11434"  # Ollama 默认地址
)

async def load_pdf():
    file_path = './examples/layout-parser-paper.pdf'
    loader = PyPDFLoader(file_path)
    pages = []
    async for page in loader.alazy_load():
        pages.append(page)

    # 后续使用方式与 OpenAIEmbeddings 完全一致
    vector_store = InMemoryVectorStore.from_documents(
        pages,
        embeddings)
    docs = vector_store.similarity_search("What is LayoutParser?", k=2)
    for doc in docs:
        print(f'Page {doc.metadata["page"]}: {doc.page_content[:300]}\n')

asyncio.run(load_pdf())
